import os
import re
import random

from tqdm      import tqdm
from evaluator import Evaluator

BASE_URL   = "http://36.111.148.95:8000/v1"
API_KEY    = "123456"


class ChatGLM3Evaluator(Evaluator):
    def __init__(self, choices, model_name, k=-1, logger=None):
        """
        @brief 构造ChatGLM3的评估器
        
        @param choices:    候选答案列表
        @param model_name: 模型名称
        @param k:          few-shot的提示数量
        """
        super().__init__(choices, model_name, base_url=BASE_URL,
                        api_key=API_KEY, k=k, logger=logger)


    def eval(self, subject, val_df, dev_df=None, few_shot=False, cot=False, save_dir=None):
        correct_num = 0
        result      = []
        score       = []

        history = ""
        if few_shot:
            history += self.gen_few_shot(subject, dev_df, cot=cot)   # 生成few-shot内容

        target = list(val_df['answer'])
        prompt = f"以下是中国关于{subject}考试的单项选择题，请选出其中的正确答案。\n\n"
        for idx, row in tqdm(val_df.iterrows(), total=len(val_df), disable=True):
            question = self.format(row, with_answer=False)

            if few_shot:
                prompt_shot = "以下是一些例子可以参考: \n"
                resp = self.client.chat.completions.create(
                        model    = self.modle_name,
                        messages = [
                            {"role": "system", "content":  f"你是{subject}学科领域的专家."},
                            {"role": "user",   "content":  prompt + prompt_shot + history + question}
                        ]
                    ).choices[0].message.content
            else:
                resp = self.client.chat.completions.create(
                        model    = self.modle_name,
                        messages = [
                            {"role": "system", "content": f"你是{subject}学科领域的专家."},
                            {"role": "user",  "content":  prompt + question}
                        ]
                    ).choices[0].message.content.strip()
            model_choice, direct_extract = self.extract_cot_answer(row, resp)

            if model_choice not in self.choices:
                self.logger.warning(f"问题{idx}模型答案{model_choice}不在参考选项内!")
                continue

            correct = 0
            if model_choice == target[idx]:
                correct_num += 1
                correct      = 1

            if save_dir:
                if few_shot:
                    result.append(resp)
                score.append(correct)
        correct_ratio = 100 * correct_num / len(target)

        if save_dir:
            if few_shot:
                val_df['model_output'] = result
            val_df['correctness'] = score
            val_df.to_csv(os.path.join(save_dir, f'{subject}_test.csv'))
        return correct_ratio

            

    def gen_few_shot(self, subject, dev_df, cot=False):
        """
        @brief 根据开发数据集合的dataframe数据生成few-shot的数据
        
        @subject 学科主题
        @dev_df  开发数据集合的dataframe数据

        @return  一些提示的问题和答案
        """
        if self.k < 1:
            return ""
        
        dev_num  = dev_df.shape[0]    # 总共的案例数量
        idx_shot = []
        if self.k >= dev_num:
            idx_shot = list(range(dev_num))
        else:
            idx_shot = random.sample(range(dev_num), self.k)
            
        
        message = ""
        for idx in idx_shot:
            msg = self.format(dev_df.iloc[idx, :], with_answer=True, cot=cot)
            message += msg
            
        return message

        
    def format(self, line, with_answer=True, cot=False, prompt=""):
        """
        @brief 格式化pandas的一行数据为指定的格式，如果带有answer表示这是个shot, 如没有answer表示这行数据是个问题.

        @line         pandas的一行数据
        @with_answer  改行数据是否包含答案
        @cot          是否使用思维链模式提供答案
        @prompt       额外的提示语句

        @return      
            example   格式化的问题
            answer    格式化的答案(如果line内有答案的话)
        """
        example = prompt + "\n" + line['question']
        for choice in self.choices:
            example += f"\n {choice}. {line[choice]}"
        example += "\n答案: "
        
        if with_answer:
            if cot:           # 通过思维链的思路将解题思路提供给模型
                answer = "让我们一步一步思考, \n" + line["explanation"] + f"\n所以答案是{line['answer']}."
            else:
                answer = line['answer']
            return example + "\n" + answer
        return example

    
    def extract_cot_answer(self, line, gen_answer):
        m = re.findall(r'所以答案是(.+?)。', gen_answer, re.M)
        if len(m) > 0 and m[-1] in self.choices:
            return m[-1], True
        answer_patterns = [
            r'([ABCD])是正确的',
            r'选项([ABCD])正确',
            r'答案为([ABCD])',
            r'答案是([ABCD])',
            r'答案([ABCD])',
            r'选择([ABCD])',
            r'答案：([ABCD])',
            r'选择答案([ABCD])'
        ]
        # RE extraction
        for answer_pattern in answer_patterns:
            m = re.search(answer_pattern, gen_answer, re.M)
            if m:
                answer = m.group(1)
                return answer, False
        # only containing one choice-character
        m = re.findall(r'[ABCD]', gen_answer, re.M)
        if len(m) == 1:
            answer = m[0]
            return answer, False
        answer_word_counter = 0
        # only containing one choice-context
        for c in self.choices:
            if str(line[f'{c}']) in gen_answer:
                answer = c
                answer_word_counter += 1
        if answer_word_counter == 1:
            return answer, False
        return '-', False
